# encoding=utf-8
import json
import os
import cv2
import numpy as np
from PIL import Image
from seg_cuda.CudaMatrixTools import MatBuilder
from seg_system.application_config import ApplicationConfig
from seg_system.Vascular.service.VascularToolsForBatch.VascularEachProcessor.VascularForBatchBase import \
    VascularForBatchBase


class CNBDForBatch(VascularForBatchBase):
    """CNBDForBatch:
        - 本算法对于并行度的支持不佳, 主要尝试使用 c++/openMP/thread的方案进行优化
        50张图的情况下, 默认开启Intersections并行时， 四种情况统计如下：
            - batch, stride, thread = 3, 15, 2 ( 测试环境 )
                - use_cxx, one_by_one, time_cost:
                - False  , False     , 180.58850858500227
                - False  , True      , 95.86277513900131
                - True   , False     , 14.0595488809995
                - True   , True      , 95.04049996100002
            - 精度损失状态(C++和python)
                - True   , False     , 0.863912080000091
            - 总结，python的多线程问题，需要使用语言扩展, 处理较大的矩阵的时候，应该选择其他方案
            - 调用语言扩展的时候，注意合并需要传输的数据，减小传输带来的大量消耗

        note:
            - Intersections部分需要重写为c++
    """

    def process_with_python(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple,
                            stride: int, **kwargs):
        """尽量保证和VascularTools/CNBD一致
        """
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        if not one_by_one:
            refine_output = self.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 这部分没改好
                img111 = self.getSkeletonIntersection(XHout, ApplicationConfig.SystemConfig.VASCULAR_USE_WORKER)
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list

    def process_with_cxx(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple, stride: int,
                         **kwargs):
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        if not one_by_one:
            refine_output = self.cxxPyVascular.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 这部分没改好, C++和python有精度差,不知道为什么
                # img111_ = self.getSkeletonIntersection(XHout, ApplicationConfig.SystemConfig.VASCULAR_USE_WORKER)
                img111 = self.cxxPyVascular.GetSkeletonIntersection(XHout.copy())
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list

    def batch_save(self, file_name: list, save_path: str, process_output, **kwargs):
        refine_output, branch_output, seg_output, CNBD_list = process_output

        # 颜色调整
        alpha = kwargs.get('alpha', 0.9)
        beta = 1.0 - alpha

        # 位置信息
        ori_save_path = kwargs.get("ori_save_path")
        brach_tmp_save_path = kwargs.get("brach_tmp_save_path")

        refine_save_path = kwargs.get("refine_save_path")
        branch_save_path = kwargs.get("branch_save_path")
        seg_save_path = kwargs.get("seg_save_path")

        for e_n, e_r, e_b, e_s in zip(file_name, refine_output, branch_output, seg_output):
            e_r_p = os.path.join(refine_save_path, e_n)
            e_b_p = os.path.join(branch_save_path, e_n)
            e_s_p = os.path.join(seg_save_path, e_n)

            cv2.imwrite(e_r_p, e_r)
            cv2.imwrite(e_b_p, e_b)
            cv2.imwrite(e_s_p, e_s)

            # 2023.1.13合并brach的输出到原图
            ori_file = cv2.imread(os.path.join(ori_save_path, e_n))
            no_zero_index = np.nonzero(e_b)
            ori_file[no_zero_index[0], no_zero_index[1], :] = \
                ori_file[no_zero_index[0], no_zero_index[1], :] * beta + \
                alpha * e_b[no_zero_index[0], no_zero_index[1], :]

            e_b_t_p = os.path.join(brach_tmp_save_path, e_n)
            cv2.imwrite(e_b_t_p, ori_file)

            assert os.path.exists(e_r_p)
            assert os.path.exists(e_b_p)
            assert os.path.exists(e_s_p)
            assert os.path.exists(e_b_t_p)

        return CNBD_list


class CNBDNoneBorderForBatch(CNBDForBatch):
    def process_with_python(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple,
                            stride: int, **kwargs):
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        if not one_by_one:
            refine_output = self.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 2022.9.13新增的去除边框部分
                d = ApplicationConfig.SystemConfig.VASCULAR_DELETE_STRIDE
                XHout_mask = np.ones(XHout.shape, np.bool)
                XHout_mask[d: XHout.shape[0] - d, d: XHout_mask.shape[1] - d] = 0
                XHout[XHout_mask] = 0

                img111 = self.getSkeletonIntersection(XHout, ApplicationConfig.SystemConfig.VASCULAR_USE_WORKER)
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list

    def process_with_cxx(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple, stride: int,
                         **kwargs):
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        if not one_by_one:
            refine_output = self.cxxPyVascular.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 2022.9.13新增的去除边框部分
                d = ApplicationConfig.SystemConfig.VASCULAR_DELETE_STRIDE
                XHout_mask = np.ones(XHout.shape, np.bool)
                XHout_mask[d: XHout.shape[0] - d, d: XHout_mask.shape[1] - d] = 0
                XHout[XHout_mask] = 0

                img111 = self.cxxPyVascular.GetSkeletonIntersection(XHout.copy())
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list

    def batch_save(self, file_name: list, save_path: str, process_output, **kwargs):
        return super().batch_save(file_name, save_path, process_output, **kwargs)


class CNBDNoneBorderWithPointForBatch(CNBDForBatch):
    """
        Date: 2023.5.9
        Author: tacom
        Description:
          根据需求保存每个图像,神经交叉点的坐标信息
          整合神经交叉点的信息到excel的下载表格中
          预计格式如下:
            文件名  PointX PointY
            aaaa      1      2
                      3      4
            bbbb      1      2
                      3      4
          预计表格分页名称:
            中文: 神经交叉坐标
            英文: Neural chiasmatic coordinates
          实现方案就是在增加vascular相关的temp文件夹,暂时存储方便解析的json数据到txt
          在download_excel接口中组合信息到表格中

         由于需要修改两个处理过程以及数据保存的流程, 为保持备份, 所以新建子类
    """

    def process_with_python(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple,
                            stride: int, **kwargs):
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        CNBD_point_list = []
        if not one_by_one:
            refine_output = self.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 2022.9.13新增的去除边框部分
                d = ApplicationConfig.SystemConfig.VASCULAR_DELETE_STRIDE
                XHout_mask = np.ones(XHout.shape, np.bool)
                XHout_mask[d: XHout.shape[0] - d, d: XHout_mask.shape[1] - d] = 0
                XHout[XHout_mask] = 0

                img111 = self.getSkeletonIntersection(XHout, ApplicationConfig.SystemConfig.VASCULAR_USE_WORKER)
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                each_point_list = {
                    'x': [],
                    'y': []
                }
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开
                    each_point_list['x'].append(int(data[0]))
                    each_point_list['y'].append(int(data[1]))
                CNBD_point_list.append(each_point_list)

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list, CNBD_point_list

    def process_with_cxx(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple, stride: int,
                         **kwargs):
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        CNBD_point_list = []
        if not one_by_one:
            refine_output = self.cxxPyVascular.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 2022.9.13新增的去除边框部分
                d = ApplicationConfig.SystemConfig.VASCULAR_DELETE_STRIDE
                XHout_mask = np.ones(XHout.shape, np.bool)
                XHout_mask[d: XHout.shape[0] - d, d: XHout_mask.shape[1] - d] = 0
                XHout[XHout_mask] = 0

                img111 = self.cxxPyVascular.GetSkeletonIntersection(XHout.copy())
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                each_point_list = {
                    'x': [],
                    'y': []
                }
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开
                    each_point_list['x'].append(int(data[0]))
                    each_point_list['y'].append(int(data[1]))
                CNBD_point_list.append(each_point_list)

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list, CNBD_point_list

    def batch_save(self, file_name: list, save_path: str, process_output, **kwargs):
        refine_output, branch_output, seg_output, CNBD_list, CNBD_point_list = process_output

        # 颜色调整
        alpha = kwargs.get('alpha', 0.9)
        beta = 1.0 - alpha

        # 位置信息
        ori_save_path = kwargs.get("ori_save_path")
        brach_tmp_save_path = kwargs.get("brach_tmp_save_path")

        refine_save_path = kwargs.get("refine_save_path")
        branch_save_path = kwargs.get("branch_save_path")
        seg_save_path = kwargs.get("seg_save_path")

        # 2023.5.13增加的神经交叉节点存储位置
        cnbd_point_path = kwargs.get("cnbd_point_path")

        for e_n, e_r, e_b, e_s, e_p in zip(file_name, refine_output, branch_output, seg_output, CNBD_point_list):
            e_r_p = os.path.join(refine_save_path, e_n)
            e_b_p = os.path.join(branch_save_path, e_n)
            e_s_p = os.path.join(seg_save_path, e_n)
            e_p_p = os.path.join(cnbd_point_path, e_n + ".txt")

            cv2.imwrite(e_r_p, e_r)
            cv2.imwrite(e_b_p, e_b)
            cv2.imwrite(e_s_p, e_s)

            # 2023.1.13合并brach的输出到原图
            ori_file = cv2.imread(os.path.join(ori_save_path, e_n))
            no_zero_index = np.nonzero(e_b)
            ori_file[no_zero_index[0], no_zero_index[1], :] = \
                ori_file[no_zero_index[0], no_zero_index[1], :] * beta + \
                alpha * e_b[no_zero_index[0], no_zero_index[1], :]

            # 2023.5.13写入交叉点list数据到文件中
            with open(e_p_p, 'w') as f:
                e_p = json.dumps(e_p)
                f.write(e_p)

            e_b_t_p = os.path.join(brach_tmp_save_path, e_n)
            cv2.imwrite(e_b_t_p, ori_file)

            assert os.path.exists(e_r_p)
            assert os.path.exists(e_b_p)
            assert os.path.exists(e_s_p)
            assert os.path.exists(e_p_p)
            assert os.path.exists(e_b_t_p)

        return CNBD_list
